{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|default_exp docments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Docments\n",
    "\n",
    "> Document parameters using comments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "from __future__ import annotations\n",
    "\n",
    "import re,ast\n",
    "from tokenize import tokenize,COMMENT\n",
    "from ast import parse,FunctionDef,AsyncFunctionDef,AnnAssign\n",
    "from io import BytesIO\n",
    "from textwrap import dedent\n",
    "from types import SimpleNamespace\n",
    "from inspect import getsource,isfunction,ismethod,isclass,signature,Parameter\n",
    "from dataclasses import dataclass, is_dataclass\n",
    "from fastcore.utils import *\n",
    "from fastcore.meta import delegates\n",
    "from fastcore import docscrape\n",
    "from inspect import isclass,getdoc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`docments` provides programmatic access to comments in function parameters and return types. It can be used to create more developer-friendly documentation, CLI, etc tools."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Without docments, if you want to document your parameters, you have to repeat param names in docstrings, since they're already in the function signature. The parameters have to be kept synchronized in the two places as you change your code. Readers of your code have to look back and forth between two places to understand what's happening. So it's more work for you, and for your users.\n",
    "\n",
    "Furthermore, to have parameter documentation formatted nicely without docments, you have to use special magic docstring formatting, often with [odd quirks](https://stackoverflow.com/questions/62167540/why-do-definitions-have-a-space-before-the-colon-in-numpy-docstring-sections), which is a pain to create and maintain, and awkward to read in code. For instance, using [numpy-style documentation](https://numpydoc.readthedocs.io/en/latest/format.html):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_np(a:int, b:int=0)->int:\n",
    "    \"\"\"The sum of two numbers.\n",
    "    \n",
    "    Used to demonstrate numpy-style docstrings.\n",
    "\n",
    "Parameters\n",
    "----------\n",
    "a : int\n",
    "    the 1st number to add\n",
    "b : int\n",
    "    the 2nd number to add (default: 0)\n",
    "\n",
    "Returns\n",
    "-------\n",
    "int\n",
    "    the result of adding `a` to `b`\"\"\"\n",
    "    return a+b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By comparison, here's the same thing using docments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add(\n",
    "    a:int, # the 1st number to add\n",
    "    b=0,   # the 2nd number to add\n",
    ")->int:    # the result of adding `a` to `b`\n",
    "    \"The sum of two numbers.\"\n",
    "    return a+b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numpy docstring helper functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`docments` also supports numpy-style docstrings, or a mix or numpy-style and docments parameter documentation. The functions in this section help get and parse this information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def docstring(sym):\n",
    "    \"Get docstring for `sym` for functions ad classes\"\n",
    "    if isinstance(sym, str): return sym\n",
    "    res = getdoc(sym)\n",
    "    if not res and isclass(sym): res = getdoc(sym.__init__)\n",
    "    return res or \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(docstring(add), \"The sum of two numbers.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def parse_docstring(sym):\n",
    "    \"Parse a numpy-style docstring in `sym`\"\n",
    "    docs = docstring(sym)\n",
    "    return AttrDict(**docscrape.NumpyDocString(docstring(sym)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parse_docstring(add_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def isdataclass(s):\n",
    "    \"Check if `s` is a dataclass but not a dataclass' instance\"\n",
    "    return is_dataclass(s) and isclass(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def get_dataclass_source(s):\n",
    "    \"Get source code for dataclass `s`\"\n",
    "    return getsource(s) if not getattr(s, \"__module__\") == '__main__' else \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def get_source(s):\n",
    "    \"Get source code for string, function object or dataclass `s`\"\n",
    "    return getsource(s) if isfunction(s) or ismethod(s) else get_dataclass_source(s) if isdataclass(s) else s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _parses(s):\n",
    "    \"Parse Python code in string, function object or dataclass `s`\"\n",
    "    return parse(dedent(get_source(s)))\n",
    "\n",
    "def _tokens(s):\n",
    "    \"Tokenize Python code in string or function object `s`\"\n",
    "    s = get_source(s)\n",
    "    return tokenize(BytesIO(s.encode('utf-8')).readline)\n",
    "\n",
    "_clean_re = re.compile(r'^\\s*#(.*)\\s*$')\n",
    "def _clean_comment(s):\n",
    "    res = _clean_re.findall(s)\n",
    "    return res[0] if res else None\n",
    "\n",
    "def _param_locs(s, returns=True):\n",
    "    \"`dict` of parameter line numbers to names\"\n",
    "    body = _parses(s).body\n",
    "    if len(body)==1: #or not isinstance(body[0], FunctionDef): return None\n",
    "        defn = body[0]\n",
    "        if isinstance(defn, (FunctionDef, AsyncFunctionDef)):\n",
    "            res = {arg.lineno:arg.arg for arg in defn.args.args}\n",
    "            if returns and defn.returns: res[defn.returns.lineno] = 'return'\n",
    "            return res\n",
    "        elif isdataclass(s):\n",
    "            res = {arg.lineno:arg.target.id for arg in defn.body if isinstance(arg, AnnAssign)}\n",
    "            return res\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "empty = Parameter.empty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _get_comment(line, arg, comments, parms):\n",
    "    if line in comments: return comments[line].strip()\n",
    "    line -= 1\n",
    "    res = []\n",
    "    while line and line in comments and line not in parms:\n",
    "        res.append(comments[line])\n",
    "        line -= 1\n",
    "    return dedent('\\n'.join(reversed(res))) if res else None\n",
    "\n",
    "def _get_full(anno, name, default, docs):\n",
    "    if anno==empty and default!=empty: anno = type(default)\n",
    "    return AttrDict(docment=docs.get(name), anno=anno, default=default)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _merge_doc(dm, npdoc):\n",
    "    if not npdoc: return dm\n",
    "    if not dm.anno or dm.anno==empty: dm.anno = npdoc.type\n",
    "    if not dm.docment: dm.docment = '\\n'.join(npdoc.desc)\n",
    "    return dm\n",
    "\n",
    "def _merge_docs(dms, npdocs):\n",
    "    npparams = npdocs['Parameters']\n",
    "    params = {nm:_merge_doc(dm,npparams.get(nm,None)) for nm,dm in dms.items()}\n",
    "    if 'return' in dms: params['return'] = _merge_doc(dms['return'], npdocs['Returns'])\n",
    "    return params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _get_property_name(p):\n",
    "    \"Get the name of property `p`\"\n",
    "    if hasattr(p, 'fget'):\n",
    "        return p.fget.func.__qualname__ if hasattr(p.fget, 'func') else p.fget.__qualname__\n",
    "    else: return next(iter(re.findall(r'\\'(.*)\\'', str(p)))).split('.')[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def get_name(obj):\n",
    "    \"Get the name of `obj`\"\n",
    "    if hasattr(obj, '__name__'):       return obj.__name__\n",
    "    elif getattr(obj, '_name', False): return obj._name\n",
    "    elif hasattr(obj,'__origin__'):    return str(obj.__origin__).split('.')[-1] #for types\n",
    "    elif type(obj)==property:          return _get_property_name(obj)\n",
    "    else:                              return str(obj).split('.')[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(get_name(in_ipython), 'in_ipython')\n",
    "test_eq(get_name(L.map), 'map')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def qual_name(obj):\n",
    "    \"Get the qualified name of `obj`\"\n",
    "    if hasattr(obj,'__qualname__'): return obj.__qualname__\n",
    "    if ismethod(obj):       return f\"{get_name(obj.__self__)}.{get_name(fn)}\"\n",
    "    return get_name(obj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert qual_name(docscrape) == 'fastcore.docscrape'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Docments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _docments(s, returns=True, eval_str=False):\n",
    "    \"`dict` of parameter names to 'docment-style' comments in function or string `s`\"\n",
    "    nps = parse_docstring(s)\n",
    "    if isclass(s) and not is_dataclass(s): s = s.__init__ # Constructor for a class\n",
    "    comments = {o.start[0]:_clean_comment(o.string) for o in _tokens(s) if o.type==COMMENT}\n",
    "    parms = _param_locs(s, returns=returns) or {}\n",
    "    docs = {arg:_get_comment(line, arg, comments, parms) for line,arg in parms.items()}\n",
    "\n",
    "    sig = signature_ex(s, True)\n",
    "    res = {arg:_get_full(p.annotation, p.name, p.default, docs) for arg,p in sig.parameters.items()}\n",
    "    if returns: res['return'] = _get_full(sig.return_annotation, 'return', empty, docs)\n",
    "    res = _merge_docs(res, nps)\n",
    "    if eval_str:\n",
    "        hints = type_hints(s)\n",
    "        for k,v in res.items():\n",
    "            if k in hints: v['anno'] = hints.get(k)\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@delegates(_docments)\n",
    "def docments(elt, full=False, **kwargs):\n",
    "    \"Generates a `docment`\"\n",
    "    r = {}\n",
    "    params = set(signature(elt).parameters)\n",
    "    params.add('return')\n",
    "\n",
    "    def _update_docments(f, r):\n",
    "        if hasattr(f, '__delwrap__'): _update_docments(f.__delwrap__, r)\n",
    "        r.update({k:v for k,v in _docments(f, **kwargs).items() if k in params\n",
    "                  and (v.get('docment', None) or not nested_idx(r, k, 'docment'))})\n",
    "\n",
    "    _update_docments(elt, r)\n",
    "    if not full: r = {k:v['docment'] for k,v in r.items()}\n",
    "    return AttrDict(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The returned `dict` has parameter names as keys, docments as values. The return value comment appears in the `return`, unless `returns=False`. Using the `add` definition above, we get:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{ 'a': 'the 1st number to add',\n",
       "  'b': 'the 2nd number to add',\n",
       "  'return': 'the result of adding `a` to `b`'}\n",
       "```"
      ],
      "text/plain": [
       "{'a': 'the 1st number to add',\n",
       " 'b': 'the 2nd number to add',\n",
       " 'return': 'the result of adding `a` to `b`'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def add(\n",
    "    a:int, # the 1st number to add\n",
    "    b=0,   # the 2nd number to add\n",
    ")->int:    # the result of adding `a` to `b`\n",
    "    \"The sum of two numbers.\"\n",
    "    return a+b\n",
    "\n",
    "docments(add)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you pass `full=True`, the values are `dict` of defaults, types, and docments as values. Note that the type annotation is inferred from the default value, if the annotation is empty and a default is supplied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{ 'a': { 'anno': <class 'int'>,\n",
       "         'default': <class 'inspect._empty'>,\n",
       "         'docment': 'the 1st number to add'},\n",
       "  'b': { 'anno': <class 'int'>,\n",
       "         'default': 0,\n",
       "         'docment': 'the 2nd number to add'},\n",
       "  'return': { 'anno': <class 'int'>,\n",
       "              'default': <class 'inspect._empty'>,\n",
       "              'docment': 'the result of adding `a` to `b`'}}\n",
       "```"
      ],
      "text/plain": [
       "{'a': {'docment': 'the 1st number to add',\n",
       "  'anno': int,\n",
       "  'default': inspect._empty},\n",
       " 'b': {'docment': 'the 2nd number to add', 'anno': int, 'default': 0},\n",
       " 'return': {'docment': 'the result of adding `a` to `b`',\n",
       "  'anno': int,\n",
       "  'default': inspect._empty}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docments(add, full=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To evaluate stringified annotations (from python 3.10), use `eval_str`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{ 'anno': <class 'int'>,\n",
       "  'default': <class 'inspect._empty'>,\n",
       "  'docment': 'the 1st number to add'}\n",
       "```"
      ],
      "text/plain": [
       "{'docment': 'the 1st number to add', 'anno': int, 'default': inspect._empty}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docments(add, full=True, eval_str=True)['a']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you need more space to document a parameter, place one or more lines of comments above the parameter, or above the return type. You can mix-and-match these docment styles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add(\n",
    "    # The first operand\n",
    "    a:int,\n",
    "    # This is the second of the operands to the *addition* operator.\n",
    "    # Note that passing a negative value here is the equivalent of the *subtraction* operator.\n",
    "    b:int,\n",
    ")->int: # The result is calculated using Python's builtin `+` operator.\n",
    "    \"Add `a` to `b`\"\n",
    "    return a+b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{ 'a': 'The first operand',\n",
       "  'b': 'This is the second of the operands to the *addition* operator.\\n'\n",
       "       'Note that passing a negative value here is the equivalent of the '\n",
       "       '*subtraction* operator.',\n",
       "  'return': \"The result is calculated using Python's builtin `+` operator.\"}\n",
       "```"
      ],
      "text/plain": [
       "{'a': 'The first operand',\n",
       " 'b': 'This is the second of the operands to the *addition* operator.\\nNote that passing a negative value here is the equivalent of the *subtraction* operator.',\n",
       " 'return': \"The result is calculated using Python's builtin `+` operator.\"}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docments(add)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Docments works with async functions, too:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def add_async(\n",
    "    # The first operand\n",
    "    a:int,\n",
    "    # This is the second of the operands to the *addition* operator.\n",
    "    # Note that passing a negative value here is the equivalent of the *subtraction* operator.\n",
    "    b:int,\n",
    ")->int: # The result is calculated using Python's builtin `+` operator.\n",
    "    \"Add `a` to `b`\"\n",
    "    return a+b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(docments(add_async), docments(add))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use docments with classes and methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Adder:\n",
    "    \"An addition calculator\"\n",
    "    def __init__(self,\n",
    "        a:int, # First operand\n",
    "        b:int, # 2nd operand\n",
    "    ): self.a,self.b = a,b\n",
    "    \n",
    "    def calculate(self\n",
    "                 )->int: # Integral result of addition operator\n",
    "        \"Add `a` to `b`\"\n",
    "        return a+b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{'a': 'First operand', 'b': '2nd operand', 'return': None}\n",
       "```"
      ],
      "text/plain": [
       "{'a': 'First operand', 'b': '2nd operand', 'return': None}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docments(Adder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{'return': 'Integral result of addition operator', 'self': None}\n",
       "```"
      ],
      "text/plain": [
       "{'self': None, 'return': 'Integral result of addition operator'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docments(Adder.calculate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "docments can also be extracted from numpy-style docstrings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The sum of two numbers.\n",
      "    \n",
      "    Used to demonstrate numpy-style docstrings.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "a : int\n",
      "    the 1st number to add\n",
      "b : int\n",
      "    the 2nd number to add (default: 0)\n",
      "\n",
      "Returns\n",
      "-------\n",
      "int\n",
      "    the result of adding `a` to `b`\n"
     ]
    }
   ],
   "source": [
    "print(add_np.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{ 'a': 'the 1st number to add',\n",
       "  'b': 'the 2nd number to add (default: 0)',\n",
       "  'return': 'the result of adding `a` to `b`'}\n",
       "```"
      ],
      "text/plain": [
       "{'a': 'the 1st number to add',\n",
       " 'b': 'the 2nd number to add (default: 0)',\n",
       " 'return': 'the result of adding `a` to `b`'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docments(add_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can even mix and match docments and numpy parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_mixed(a:int, # the first number to add\n",
    "              b\n",
    "             )->int: # the result\n",
    "    \"\"\"The sum of two numbers.\n",
    "\n",
    "Parameters\n",
    "----------\n",
    "b : int\n",
    "    the 2nd number to add (default: 0)\"\"\"\n",
    "    return a+b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{ 'a': { 'anno': <class 'int'>,\n",
       "         'default': <class 'inspect._empty'>,\n",
       "         'docment': 'the first number to add'},\n",
       "  'b': { 'anno': 'int',\n",
       "         'default': <class 'inspect._empty'>,\n",
       "         'docment': 'the 2nd number to add (default: 0)'},\n",
       "  'return': { 'anno': <class 'int'>,\n",
       "              'default': <class 'inspect._empty'>,\n",
       "              'docment': 'the result'}}\n",
       "```"
      ],
      "text/plain": [
       "{'a': {'docment': 'the first number to add',\n",
       "  'anno': int,\n",
       "  'default': inspect._empty},\n",
       " 'b': {'docment': 'the 2nd number to add (default: 0)',\n",
       "  'anno': 'int',\n",
       "  'default': inspect._empty},\n",
       " 'return': {'docment': 'the result', 'anno': int, 'default': inspect._empty}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docments(add_mixed, full=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use docments with dataclasses, however if the class was defined in online notebook, docments will not contain parameters' comments. This is because the source code is not available in the notebook. After converting the notebook to a module, the docments will be available. Thus, documentation will have correct parameters' comments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "class _F:\n",
    "    @classmethod\n",
    "    def class_method(cls, \n",
    "                     foo:str, # docment for parameter foo\n",
    "                     ):...\n",
    "    \n",
    "test_eq(docments(_F.class_method), {'foo': 'docment for parameter foo', 'return': None})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Docments even works with `delegates`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastcore.meta import delegates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{'a': 'First', 'b': 'Second', 'return': None}\n",
       "```"
      ],
      "text/plain": [
       "{'a': 'First', 'return': None, 'b': 'Second'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _a(a:int=2): return a # First\n",
    "\n",
    "@delegates(_a)\n",
    "def _b(b:str, **kwargs): return b, (_a(**kwargs)) # Second\n",
    "\n",
    "docments(_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "def _c(b:str, # Second\n",
    "       a:int=2): return b, a # Third\n",
    "\n",
    "@delegates(_c)\n",
    "def _d(c:int, # First\n",
    "       b:str, **kwargs\n",
    "      )->int:\n",
    "    return c, _c(b, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(docments(_c, full=True)['b']['docment'],'Second')\n",
    "test_eq(docments(_d, full=True)['b']['docment'],'Second')\n",
    "_argset = {'a', 'b', 'c', 'return'}\n",
    "test_eq(docments(_d, full=True).keys() & _argset, _argset) # _d has the args a,b,c and return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract docstrings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _get_params(node):\n",
    "    params = [a.arg for a in node.args.args]\n",
    "    if node.args.vararg: params.append(f\"*{node.args.vararg.arg}\")\n",
    "    if node.args.kwarg: params.append(f\"**{node.args.kwarg.arg}\")\n",
    "    return \", \".join(params)\n",
    "\n",
    "class _DocstringExtractor(ast.NodeVisitor):\n",
    "    def __init__(self): self.docstrings,self.cls,self.cls_init = {},None,None\n",
    "\n",
    "    def visit_FunctionDef(self, node):\n",
    "        name = node.name\n",
    "        if name == '__init__':\n",
    "            self.cls_init = node\n",
    "            return\n",
    "        elif name.startswith('_'): return\n",
    "        elif self.cls: name = f\"{self.cls}.{node.name}\"\n",
    "        docs = ast.get_docstring(node)\n",
    "        params = _get_params(node)\n",
    "        if docs: self.docstrings[name] = (docs, params)\n",
    "        self.generic_visit(node)\n",
    "\n",
    "    def visit_ClassDef(self, node):\n",
    "        self.cls,self.cls_init = node.name,None\n",
    "        docs = ast.get_docstring(node)\n",
    "        if docs: self.docstrings[node.name] = ()\n",
    "        self.generic_visit(node)\n",
    "        if not docs and self.cls_init: docs = ast.get_docstring(self.cls_init)\n",
    "        params = _get_params(self.cls_init) if self.cls_init else \"\"\n",
    "        if docs: self.docstrings[node.name] = (docs, params)\n",
    "        self.cls,self.cls_init = None,None\n",
    "\n",
    "    def visit_Module(self, node):\n",
    "        module_doc = ast.get_docstring(node)\n",
    "        if module_doc: self.docstrings['_module'] = (module_doc, \"\")\n",
    "        self.generic_visit(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def extract_docstrings(code):\n",
    "    \"Create a dict from function/class/method names to tuples of docstrings and param lists\"\n",
    "    extractor = _DocstringExtractor()\n",
    "    extractor.visit(ast.parse(code))\n",
    "    return extractor.docstrings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_code = \"\"\"\n",
    "\"This is a module.\"\n",
    "\n",
    "def top_func(a, b, *args, **kw):\n",
    "    \"This is top-level.\"\n",
    "    pass\n",
    "\n",
    "class SampleClass:\n",
    "    \"This is a class.\"\n",
    "\n",
    "    def __init__(self, x, y):\n",
    "        \"Constructor for SampleClass.\"\n",
    "        pass\n",
    "\n",
    "    def method1(self, param1):\n",
    "        \"This is method1.\"\n",
    "        pass\n",
    "\n",
    "    def _private_method(self):\n",
    "        \"This should not be included.\"\n",
    "        pass\n",
    "\n",
    "class AnotherClass:\n",
    "    def __init__(self, a, b):\n",
    "        \"This class has no separate docstring.\"\n",
    "        pass\"\"\"\n",
    "\n",
    "exp = {'_module': ('This is a module.', ''),\n",
    "       'top_func': ('This is top-level.', 'a, b, *args, **kw'),\n",
    "       'SampleClass': ('This is a class.', 'self, x, y'),\n",
    "       'SampleClass.method1': ('This is method1.', 'self, param1'),\n",
    "       'AnotherClass': ('This class has no separate docstring.', 'self, a, b')}\n",
    "test_eq(extract_docstrings(sample_code), exp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export -"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
